"""
https://github.com/offu/WeRoBot/issues/486
腾讯的限制
[BUG] handle里面函数处理超过5秒会连续调用3次
"""

import werobot
import datetime
import sys
import os
import time
import redis
import pymongo as pm
import threading
from langchain.llms import OpenAI
from neo4j import GraphDatabase

sys.path.insert(0, '/home/yunpeng/code')
from PyCmpltrtok.common import sep, to_decode, to_int, encode_set, decode_set, list2str_with_limit
from PyCmpltrtok.util_mongo import delete_many_by_user, get_history, enqueue
from common import MONGODB_NAME, VALUE, KEY, IO_PREFIX, CHAT_ONLY, CHAT_AND_DIAG, CHAT_MODES, RESULT_PREFIX, SWITCH_COMMAND_SET, CHAT_PROMPT_DICT, SWITCH_DONE_DICT, PROMPT_MED, GENERATING, CHAT_MODE_NEW_DICT
from fastapi_interface import text_gen
from xlangchain_gen import analyze_symptom4user
from python_nlp.xfaiss.flask_faiss_on_mongo_test import search
from python_nlp.kg.neo4j.seek_data_with_embed.xdiagnose import get_symptom, diagnose

os.environ['R_HOST'] = '127.0.0.1'
os.environ['R_PORT'] = '6379'
os.environ['R_PWD'] = 'lgdz4qEdt/ezElyQnXFYXB80iM3OxEbAWRjMFPcIXH5ni6eQ8QOlfp7G7gvV1svPu2Bv7v'
os.environ['M_HOST'] = '127.0.0.1'
os.environ['M_PORT'] = '27017'
os.environ['M_USER'] = 'root'
os.environ['M_PWD'] = 'p1983mdA1-ei'
os.environ['W_PORT'] = '8081'  # wechat server port
os.environ['O_PORT'] = '30000'  # openai api port
os.environ['N_PORT'] = '7683'  # neo4j port
os.environ['S_PORT'] = '7700'  # FAISS search service port

S_PORT = int(os.environ['S_PORT'])
DIS_THRESH = 20.0  # 搜索时距离上限
WORDS_LIMIT = 4  # 搜索后最多用这么多症状
THRESH_S_DIAG = 45  # 诊断超过这么多秒就不作数了
CHAR_LIMIT = 450

model = OpenAI(
    streaming=True,
    verbose=True,
    callbacks=[],
    openai_api_key="token1",
    openai_api_base=f"http://127.0.0.1:{os.environ['O_PORT']}/v1",
    model_name="Qwen-1_8B-Chat",
    temperature=0.0,
    openai_proxy=None,
    top_p=1.0,
)

# 连接neo4j
print('Connecting neo4j ...', flush=True)
driver = GraphDatabase.driver(**{
    "uri": f"bolt://127.0.0.1:{os.environ['N_PORT']}",
    "auth": ("neo4j", "p1983ncbdA1-ei"),
    "encrypted": False
})
print('Connected to neo4j.', flush=True)

# 连接Redis
sep('Redis')
rdb = redis.Redis(
    os.environ['R_HOST'], int(os.environ['R_PORT']), 0, socket_timeout=3,
    password=os.environ['R_PWD'],
)
rdb.get('try_it')
sep('Redis OK')

# 连接Mongodb
sep('MongoDB')
mongo = pm.MongoClient(
    os.environ['M_HOST'], int(os.environ['M_PORT']), serverSelectionTimeoutMS=3000,
    username=os.environ['M_USER'], password=os.environ['M_PWD'], 
)
mdb = mongo[MONGODB_NAME]
get_history(mdb, 'user_xxxx', limit=1)  # try it
sep('MongoDB OK')

sep('Werobot')
robot = werobot.WeRoBot(token='kx8M7eODU4HiyAmw3a')
robot.config['APP_ID'] = 'wx47577a029b258f6b'
robot.config['ENCODING_AES_KEY'] = 'LjBrnwSvqGXx89Tszleo4sMYp8yMALcTArwhPjxoZGJ'
client = robot.client
sep('Werobot OK')

CHAR_LIMIT = 450
PROMPT_SHORT = '（请继续对话，或者输入逗号（,或，）来重置对话。）'
PROMPT = '（因平台限制，内容未完全展示，输入句号（.或。）来继续展示。请继续对话，或者输入逗号（,或，）来重置对话。）'
GO_ON_SET = set(['.', '。'])
RESET_SET = set([',', '，'])
WAIT_LAST_TURN = '（请等待上一轮对话完成。）'
THRESH_S = 20  # 单例锁最多等待这么多秒

if 0:
    """
    自定义菜单 需要 微信认证
    微信认证： 账号主体为个人，无法开通微信认证
    """
    client.create_menu({
        'button': [{
            'type': 'click',
            'name': '继续',
            'key': 'go_on'
        }]
    })


    @robot.key_click('go_on')
    def go_on(message):    
        pass
    

@robot.handler
def hello(message):
    
    sep()
    
    # 判断类型
    xtype = message.type
    print(datetime.datetime.now(), xtype, flush=True)
    if 'text' != xtype:
        return '目前仅支持文本信息！' + PROMPT_SHORT
    
    # 用户标识
    xuserid = message.source
    print('userid:', type(xuserid), xuserid, flush=True)
    
    # 输入文本
    xinput = message.content.strip()
    print('input:', xinput, flush=True)
    
    # 单例限制
    if xinput not in GO_ON_SET:
        xis_in = rdb.sismember('users_now_lock', xuserid)
        ts_now = time.time_ns()
        if xis_in:
            ts = rdb.hget('username2ts', xuserid)
            ts = to_int(ts)
            if ts and ts_now - ts < THRESH_S * 1e9:
                users_count = rdb.scard('users_now')
                xoutput = WAIT_LAST_TURN + PROMPT_SHORT + f'【活跃用户{users_count}人】'
                print('OUTPUT:', xoutput, flush=True)
                return xoutput
        rdb.hset('username2ts', xuserid, ts_now)
        rdb.sadd('users_now_lock', xuserid)
        
    # mode
    mode = rdb.hget('user2mode', xuserid)
    mode = to_decode(mode, default=CHAT_ONLY)
    if mode not in CHAT_MODES:
        mode = CHAT_ONLY
    SWITCH_PROMPT = CHAT_PROMPT_DICT[mode]
    
    if xinput in SWITCH_COMMAND_SET:
        rdb.hset('user2mode', xuserid, CHAT_MODE_NEW_DICT[mode].encode('utf8'))
        # 解除单例限制
        rdb.srem('users_now_lock', xuserid)
        return SWITCH_DONE_DICT[mode]
        
    def do_return(xoutput):
        """
        闭包函数，大部分返回走这里。
        """
        users_count = rdb.scard('users_now')
        xoutput += f'【活跃用户{users_count}人】' + SWITCH_PROMPT
        print('OUTPUT:', xoutput, flush=True)
        # 解除单例限制
        rdb.srem('users_now_lock', xuserid)
        return xoutput

    def get_result():
        """
        闭包函数，拿数据
        """
        # 最后结果
        xoutput = rdb.hget('username2output', xuserid)
        
        # 当前偏移
        xoffset = rdb.hget('username2offset', xuserid)
        xoffset = to_int(xoffset)
        xoffset_ = xoffset
        # 分推理结束和未结束两种
        if xoutput is not None:  # 结束
            xoutput = xoutput.decode('utf8')
            
            # 取本次输出
            xoutput = xoutput[xoffset:]
            xlen_output = len(xoutput)
            
            # 分输出完成未完成两种
            if xlen_output <= CHAR_LIMIT:
                # 输出完成，清理username2dynamic以标志输出完成
                rdb.hdel('username2dynamic', xuserid)
                rdb.hdel('username2offset', xuserid)
                return xoutput + PROMPT_SHORT
            else:
                # 输出未完成，分段输出
                xres = xoutput[:CHAR_LIMIT]
                
        else:  # 未结束
            
            # 从动态响应里拿
            xres = rdb.hget('username2dynamic', xuserid)
            xres = to_decode(xres)
            
            # 限制一下长度
            xres = xres[xoffset:xoffset+CHAR_LIMIT]
            
        print('output:', xres, flush=True)
        # 记录偏移
        xlen = len(xres)
        rdb.hset('username2offset', xuserid, xoffset_ + xlen)
        
        return xres + PROMPT
    
    if xinput in GO_ON_SET:  # "继续显示"
        # 从动态响应里拿
        xres = rdb.hget('username2dynamic', xuserid)
        # 没有动态响应，说明输出结束
        if xres is None:
            return do_return('（已经没有需要继续输出的内容了。）' + PROMPT_SHORT)
        
        print('GO ON >>>>', flush=True)
        users_count = rdb.scard('users_now')
        xoutput = get_result() + f'【活跃用户{users_count}人】' + SWITCH_PROMPT
        print('OUTPUT:', xoutput, flush=True)
        return xoutput
    elif xinput in RESET_SET:  # “重置对话”
        delete_many_by_user(mdb, 'dialog_in', xuserid)
        delete_many_by_user(mdb, 'dialog_out', xuserid)
        rdb.hdel('user2dis', xuserid)
        return do_return('对话已经重置。')
    
    # 生成
    rdb.sadd('users_now', xuserid)
    ts = time.time_ns()
    enqueue(mdb, 'dialog_in', xuserid, ts, {
        VALUE: xinput,
        'io': 'i',
    })
    gen = text_gen(xinput, xuserid)
    
    def generate_text():
        rdb.hdel('username2output', xuserid)
        output = ''
        for x in gen:
            rdb.hset('username2dynamic', xuserid, x.encode('utf8'))
            output = x
        rdb.hset('username2output', xuserid, output.encode('utf8'))
        rdb.srem('users_now', xuserid)
        if CHAT_AND_DIAG != mode:
            ts = time.time_ns()
            enqueue(mdb, 'dialog_out', xuserid, ts, {
                VALUE: x,
                'io': 'o',
            })
            
    generate_text_th = threading.Thread(target=generate_text)
    generate_text_th.start()
    
    if CHAT_AND_DIAG != mode:
        def check_quick_reply():
            """
            闭包函数，侦测快速返回，并获取。
            """
            # 最后结果
            xoutput = rdb.hget('username2output', xuserid)
            
            # 如果还没有结束
            if xoutput is None:  # 结束
                return None
            
            # 已经结束
            xoutput = xoutput.decode('utf8')
            
            # 当前偏移
            xoffset = rdb.hget('username2offset', xuserid)
            xoffset = to_int(xoffset)
            xoffset_ = xoffset
            
            # 取本次输出
            xoutput = xoutput[xoffset:]
            xlen_output = len(xoutput)
            
            # 分输出完成未完成两种
            if xlen_output <= CHAR_LIMIT:
                # 输出完成，清理username2dynamic以标志输出完成
                rdb.hdel('username2dynamic', xuserid)
                rdb.hdel('username2offset', xuserid)
                print('output:', xoutput, flush=True)
                return xoutput + PROMPT_SHORT
            else:
                # 输出未完成，分段输出
                xres = xoutput[:CHAR_LIMIT]
                print('output:', xres, flush=True)
                
                # 记录偏移
                xlen = len(xres)
                rdb.hset('username2offset', xuserid, xoffset_ + xlen)
                
                return xres + PROMPT
        
        # 既然限制5秒，干脆等3.75秒，让大模型去生成
        for _ in range(3):
            time.sleep(1.0)
            r = check_quick_reply()
            if r is not None:
                return do_return(r)
        time.sleep(0.75)
        
        # 3.75秒后，拿结果
        xres = get_result()
        return do_return(xres)
    
    else:
        
        def analyze_disease():
            # 以删除来标志开始
            rdb.hdel('user2result', xuserid)
        
            # 当前用户的结果
            xencoded = rdb.hget('user2dis', xuserid)
            xencoded = to_decode(xencoded)
            xdis_set = set(decode_set(xencoded))
            xdis_list_ = sorted(xdis_set)
            sep()
            print(f'当前诊断：{xdis_list_}', flush=True)
            
            xsymptoms = analyze_symptom4user(mdb, xinput, xuserid, model)
            
            # 搜索输入，选出距离小于阈值的结果
            for xsym in xsymptoms:
                sep(xsym)
                xoutput = search(xsym, 8, S_PORT)
                xsents = xoutput['sentences']
                xds = xoutput['D']
                print('候选:', flush=True)
                xwords = []
                for i, s in enumerate(xsents):
                    d = xds[i]
                    print(i, s, d, d < DIS_THRESH, flush=True)
                    if d < DIS_THRESH:
                        xwords.append(s)
            xwords = sorted(set(xwords))

            # 从这些结果中挑选出是“症状”的节点名字
            xlist = get_symptom(driver, xwords, WORDS_LIMIT)
            print(f'输入了症状：{xlist}', flush=True)
            
            # 用这些症状去诊断
            xdiseases = diagnose(driver, xlist)
            if xdiseases:
                print(f'本轮诊断：{xdiseases}', flush=True)
                if not xdis_set:
                    # 第一次诊断
                    xdis_set = set(xdiseases)
                else:
                    # 后续诊断
                    xresult = xdis_set.intersection(set(xdiseases))  # 求交集
                    if len(xresult):
                        print('【有交集】', flush=True)
                    else:
                        # 没有交集求并集
                        print('【没有交集求并集】', flush=True)
                        xresult = xdis_set.union(set(xdiseases))
                    xdis_set = xresult
                    print('交并后:', xdis_set, flush=True)
                        
            # 结果            
            xdis_list = sorted(xdis_set)
            
            # 结果还是那样，就不使用
            if xdis_list_ == xdis_list:
                print('【诊断结果没有变化】', flush=True)
                rdb.hset('user2result', xuserid, '')
            else:
                # 保存用户结果
                xencoded = encode_set(xdis_set)
                rdb.hset('user2dis', xuserid, xencoded.encode('utf8'))
                
                # 返回诊断
                xoutput = RESULT_PREFIX + list2str_with_limit(xdis_list, CHAR_LIMIT)
                if len(xdis_list) > 5:
                    xoutput += '（可能性较多，请进一步描述症状来缩小范围。）'
                rdb.hset('user2result', xuserid, xoutput.encode('utf8'))
                
        analyze_th = threading.Thread(target=analyze_disease)
        analyze_th.start()
        
        time.sleep(3.75)
        user2result = rdb.hget('user2result', xuserid)
        user2result = to_decode(user2result, default=None)
        if user2result:
            xoutput = user2result + PROMPT_MED + PROMPT_SHORT
            print('OUTPUT 1:', xoutput, flush=True)
            return do_return(xoutput)
        elif user2result is None:
            xoutput = GENERATING + PROMPT_SHORT
            print('OUTPUT 2:', xoutput, flush=True)
            return xoutput
        else:
            xoutput = get_result()
            print('OUTPUT 3:', xoutput, flush=True)
            return do_return(xoutput)
        

# 让服务器监听在 0.0.0.0:80
robot.config['HOST'] = '0.0.0.0'
robot.config['PORT'] = int(os.environ['W_PORT'])
robot.run()
